package com.shuai.ai_core.service.impl;

import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.json.JSONUtil;
import com.shuai.ai_core.dto.ChatReq;
import com.shuai.ai_core.dto.UserDTO;
import com.shuai.ai_core.service.AIInteractService;
import com.shuai.common.constant.ConversationRoleEnum;
import com.shuai.common.constant.SessionTypeWeightEnum;
import com.shuai.common.constant.UserConstant;
import com.shuai.common.constant.UserLevelTokenRelatedEnum;
import com.shuai.common.dto.SessionChatRecordDTO;
import com.shuai.common.service.user.InnerSessionChatRecordService;
import com.shuai.common.service.user.InnerUserService;
import lombok.extern.slf4j.Slf4j;
import org.apache.dubbo.config.annotation.DubboReference;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.util.StopWatch;
import reactor.core.publisher.Flux;

@Service
@Slf4j
public class AIInteractServiceImpl implements AIInteractService {
    @Autowired
    private OpenAiChatClient openAiChatClient;

    @Autowired
    private StreamingChatClient streamingChatClient;

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @DubboReference
    private InnerSessionChatRecordService sessionChatRecordService;

    @DubboReference
    private InnerUserService innerUserService;


    @Override
    public ChatResponse interactCommonPrint(ChatReq chatReq, String userId) {
        // 根据用户ID查询用户信息
        String userJson = stringRedisTemplate.opsForValue().get(UserConstant.USER_ATTRIBUTE_REDIS_KEY + userId);
        UserDTO userInfo = JSONUtil.toBean(userJson, UserDTO.class);
        // 拿到用户类型，填充每次请求的token数量
        Integer level = userInfo.getLevel();
        UserLevelTokenRelatedEnum tokenRpsLimit = UserLevelTokenRelatedEnum.getById(level); // 根据用户等级获取模型回复token的上限
        UserMessage userMessage = new UserMessage(chatReq.getQuestion());
        OpenAiChatOptions chatOptions = OpenAiChatOptions.builder()
                .withMaxTokens(tokenRpsLimit.getTokens()).build();
        Prompt prompt = new Prompt(userMessage, chatOptions);
        ChatResponse chatResponse = openAiChatClient.call(prompt);
        // 拿到模型回复结果长度，计算token
        String rpsContent = chatResponse.getResult().getOutput().getContent();
        this.updateUserTokenAndSaveSessionRecord(chatReq, rpsContent, userId);
        // 返回模型响应结果
        return chatResponse;
    }

    @Override
    public Flux<String> interactStreamPrint(ChatReq chatReq, String userId) {
        String userJson = stringRedisTemplate.opsForValue().get(UserConstant.USER_ATTRIBUTE_REDIS_KEY + userId);
        UserDTO userInfo = JSONUtil.toBean(userJson, UserDTO.class);
        Integer level = userInfo.getLevel();
        UserLevelTokenRelatedEnum tokenRpsLimit = UserLevelTokenRelatedEnum.getById(level);
        UserMessage userMessage = new UserMessage(chatReq.getQuestion());
        OpenAiChatOptions chatOptions = OpenAiChatOptions.builder()
                .withMaxTokens(tokenRpsLimit.getTokens()).build();
        Prompt prompt = new Prompt(userMessage, chatOptions);
        StringBuilder builder = new StringBuilder();
        Flux<String> flux = streamingChatClient.stream(prompt)
                .flatMap(response -> {
                    String result = response.getResult().getOutput().getContent();
                    if (CharSequenceUtil.isBlank(result)) {
                        return Flux.empty();
                    }
                    builder.append(result);
                    return Flux.just(result);
                });
        return flux.doOnComplete(() -> {
            StopWatch stopWatch = new StopWatch();
            stopWatch.start();
            this.updateUserTokenAndSaveSessionRecord(chatReq, builder.toString(), userId);
            stopWatch.stop();
            log.info("远程用户服务调用耗时：{}", stopWatch.getTotalTimeMillis());
        });
    }

    private void updateUserTokenAndSaveSessionRecord(ChatReq chatReq, String rpsContent, String userId) {
        int tokenConsume = (int) Math.floor(rpsContent.length() * SessionTypeWeightEnum.NORMAL.getWeight());
        // 扣减用户token，Dubbo服务会更新用户信息缓存
        innerUserService.decreaseUserToken(tokenConsume, userId);
        // 保存会话信息
        // 用户请求
        SessionChatRecordDTO reqDTO = SessionChatRecordDTO.builder()
                .content(chatReq.getQuestion())
                .role(ConversationRoleEnum.USER.getRole())
                .tokenNum(tokenConsume)
                .sessionId(chatReq.getSessionId())
                .build();
        sessionChatRecordService.saveSessionChatRecord(reqDTO);
        // 大模型回复
        SessionChatRecordDTO rpsDTO = SessionChatRecordDTO.builder()
                .content(rpsContent)
                .role(ConversationRoleEnum.SYSTEM.getRole())
                .tokenNum(0)
                .sessionId(chatReq.getSessionId())
                .build();
        sessionChatRecordService.saveSessionChatRecord(rpsDTO);
    }
}
